package com.hzw.video.extractor

import android.media.MediaExtractor
import android.media.MediaFormat
import java.nio.ByteBuffer

/**
 *
 * @Author:         HuaZhongWei
 * @CreateDate:     2020/9/16 9:19
 *
 * @Description:
 *
 */
class MMExtractor(path: String) {

    /**音视频分离器*/
    private val extractor: MediaExtractor = MediaExtractor().also {
        it.setDataSource(path)
    }

    /**音频通道索引*/
    private var _audioTrack = -1
    val audioTrack
        get() = _audioTrack

    /**视频通道索引*/
    private var _videoTrack = -1
    val videoTrack
        get() = _videoTrack

    /**当前帧时间戳*/
    private var _currentSampleTime = 0L
    val currentSampleTime
        get() = _currentSampleTime

    /**开始解码时间点*/
    var startPos = 0L

    /**
     * 获取视频格式参数
     */
    fun getVideoFormat(): MediaFormat? {
        for (i in 0 until extractor.trackCount) {
            val format = extractor.getTrackFormat(i)
            val mime = format.getString(MediaFormat.KEY_MIME)
            if (mime != null && mime.startsWith("video/")) {
                _videoTrack = i
                break
            }
        }
        return if (_videoTrack >= 0) {
            extractor.getTrackFormat(_videoTrack)
        } else {
            null
        }
    }

    /**
     * 获取音频格式参数
     */
    fun getAudioFormat(): MediaFormat? {
        for (i in 0 until extractor.trackCount) {
            val format = extractor.getTrackFormat(i)
            val mime = format.getString(MediaFormat.KEY_MIME)
            if (mime != null && mime.startsWith("audio/")) {
                _audioTrack = i
                break
            }
        }
        return if (_audioTrack >= 0) {
            extractor.getTrackFormat(_audioTrack)
        } else {
            null
        }
    }

    /**
     * 读取视频数据
     */
    fun readBuffer(byteBuffer: ByteBuffer): Int {
        byteBuffer.clear()
        selectSourceTrack()
        val readSampleCount = extractor.readSampleData(byteBuffer, 0)
        if (readSampleCount < 0) {
            return -1
        }
        //记录当前帧的时间戳
        _currentSampleTime = extractor.sampleTime
        //进入下一帧
        extractor.advance()
        return readSampleCount
    }

    /**
     * 选择通道
     */
    private fun selectSourceTrack() {
        if (_videoTrack >= 0) {
            extractor.selectTrack(_videoTrack)
        } else if (_audioTrack >= 0) {
            extractor.selectTrack(_audioTrack)
        }
    }

    /**
     * Seek到指定位置，并返回实际帧的时间戳
     */
    fun seek(pos: Long): Long {
        // seek(pos: Long)方法，主要用于跳播，快速将数据定位到指定的播放位置，但是，由于视频中，除了I帧以外，PB帧都需要依赖其他的帧进行解码，
        // 所以，通常只能seek到I帧，但是I帧通常和指定的播放位置有一定误差，因此需要指定seek靠近哪个关键帧，有以下三种类型：
        // SEEK_TO_PREVIOUS_SYNC：跳播位置的上一个关键帧
        // SEEK_TO_NEXT_SYNC：跳播位置的下一个关键帧
        // SEEK_TO_CLOSEST_SYNC：距离跳播位置的最近的关键帧
        extractor.seekTo(pos, MediaExtractor.SEEK_TO_PREVIOUS_SYNC)
        return extractor.sampleTime
    }

    fun stop() {
        extractor.release()
    }

    fun getSampleFlag(): Int = extractor.sampleFlags

}